package com.luncheng.mcp.core;

import com.luncheng.mcp.core.factory.McpClientFactory;
import com.luncheng.mcp.core.load.McpConfigLoader;
import com.luncheng.mcp.core.model.McpClientConfig;
import io.modelcontextprotocol.client.McpSyncClient;
import jakarta.annotation.PostConstruct;
import jakarta.annotation.PreDestroy;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

import java.util.Map;
import java.util.Objects;
import java.util.concurrent.ConcurrentHashMap;

@Slf4j
@Component
public class McpClientServer {
    @Autowired()
    private McpConfigLoader configLoader;

    private static final Map<String, McpSyncClient> clientMap = new ConcurrentHashMap<>();

    @PostConstruct
    public void start() {
        try {
            log.info("正在初始化 MCP 客户端...");
            if (Objects.nonNull(configLoader)) {
                for (McpClientConfig config : configLoader.loadClientConfig()) {
                    addMcpSyncServer(config);
                }
            }
            log.info("MCP 客户端初始化完成，共加载 {} 个服务", clientMap.size());
        } catch (Exception e) {
            log.error("MCP 客户端初始化失败", e);
            throw new RuntimeException("MCP 客户端初始化失败", e);
        }
    }

    public void addMcpSyncServer(McpClientConfig mcpClientConfig) {
        log.info("正在添加 MCP 客户端: {}", mcpClientConfig.getName());
        try {
            McpSyncClient client = McpClientFactory.createClient(mcpClientConfig);
            clientMap.put(mcpClientConfig.getName(), client);
            log.info("MCP 客户端 [{}] 添加成功", mcpClientConfig.getName());
        } catch (Exception e) {
            log.error("创建 MCP 客户端 [{}] 失败", mcpClientConfig.getName(), e);
            throw new RuntimeException("创建 MCP 客户端失败: " + mcpClientConfig.getName(), e);
        }
    }

    public void removeMcpSyncServer(String name) {
        log.info("正在移除 MCP 客户端: {}", name);
        McpSyncClient mcpSyncServer = clientMap.remove(name);
        if (mcpSyncServer != null) {
            mcpSyncServer.close();
            log.info("MCP 客户端 [{}] 已关闭并移除", name);
        } else {
            log.warn("尝试移除不存在的 MCP 客户端: {}", name);
        }
    }

    @PreDestroy
    public void stop() {
        log.info("正在停止所有 MCP 客户端...");
        clientMap.forEach((name, server) -> {
            log.info("正在关闭 MCP 客户端: {}", name);
            server.close();
        });
        clientMap.clear();
        log.info("所有 MCP 客户端已停止");
    }

    public static Map<String, McpSyncClient> getClientMap() {
        return clientMap;
    }
}